#!/usr/bin/env python
# coding: utf-8

# ### 解压掌纹数据集
get_ipython().system('cd data/data148864 && unzip -q PalmBigDataBase.zip')


# ### 简单处理数据
# 对本地数据做一些预处理，主要用于随机分组，一部分用于训练，一部分用于验证。同时生产两者的标签文件

import codecs
import os
import random
import shutil
from PIL import Image

all_file_dir = 'data/data148864/PalmBigDataBase/PalmBigDataBase'

os.mkdir('data/PalmBigDataBase/')
for i in range(1,387):
    dir = os.path.join('data/PalmBigDataBase/', str(i))
    if not os.path.exists(dir):
        os.makedirs(dir)

file_list = [c for c in os.listdir(all_file_dir) if c.endswith('bmp')]
file_list.sort()
print(file_list)

import re
import cv2

for s in file_list:
    num = re.findall(r'\d+',s)
    #print(type(num[0]))

    string1 = "P_S_" + num[0]
    string2 = "P_F_" + num[0]
    if string1 in s or string2 in s:
        img = cv2.imread('data/data148864/PalmBigDataBase/PalmBigDataBase/'+s, cv2.IMREAD_GRAYSCALE)
        cv2.imwrite('data/PalmBigDataBase/'+num[0]+'/'+s, img)

file_dir = 'data/PalmBigDataBase'
file_num = []
# file_num[0] = 0
for i in range(0,387):
    file_num.append(0)
# print(os.listdir(file_dir))
for c in os.listdir(file_dir):
    num = int(c)
    # print(num)
    for i in os.listdir(file_dir+'/'+c):
        # print(i)
        if i.endswith('bmp'):
            file_num[num] += 1

# file_num.sort()
print(file_num)

from numpy import *
num = mean(file_num)
num

import os
from shutil import copy, rmtree
import random


def mk_file(file_path: str):
    if os.path.exists(file_path):
        # 如果文件夹存在，则先删除原文件夹在重新创建
        rmtree(file_path)
    os.makedirs(file_path)


def main():
    # 保证随机可复现
    random.seed(0)

    # 将数据集中10%的数据划分到验证集中
    split_rate = 0.2

    # 指向你解压后的flower_photos文件夹

    data_root = os.getcwd()
    print(data_root)
    origin_path = os.path.join(data_root, "img")
    assert os.path.exists(origin_path), "path '{}' does not exist.".format(origin_path)

    flower_class = [cla for cla in os.listdir(origin_path)
                    if os.path.isdir(os.path.join(origin_path, cla))]

    # 建立保存训练集的文件夹
    train_root = os.path.join(data_root, "train")
    mk_file(train_root)
    for cla in flower_class:
        # 建立每个类别对应的文件夹
        mk_file(os.path.join(train_root, cla))

    # 建立保存验证集的文件夹
    val_root = os.path.join(data_root, "val")
    mk_file(val_root)
    for cla in flower_class:
        # 建立每个类别对应的文件夹
        mk_file(os.path.join(val_root, cla))

    for cla in flower_class:
        cla_path = os.path.join(origin_path, cla)
        images = os.listdir(cla_path)
        num = len(images)
        # 随机采样验证集的索引
        eval_index = random.sample(images, k=int(num * split_rate))
        for index, image in enumerate(images):
            if image in eval_index:
                # 将分配至验证集中的文件复制到相应目录
                image_path = os.path.join(cla_path, image)
                new_path = os.path.join(val_root, cla)
                copy(image_path, new_path)
            else:
                # 将分配至训练集中的文件复制到相应目录
                image_path = os.path.join(cla_path, image)
                new_path = os.path.join(train_root, cla)
                copy(image_path, new_path)
            print("\r[{}] processing [{}/{}]".format(cla, index + 1, num), end="")  # processing bar
        print()

    print("processing done!")


if __name__ == '__main__':
    main()

